from yacs.config import CfgNode as CN

_C = CN()
_C.DEVELPOER_MODE = False  # 是否开启开发者模式
_C.RANDOM_SEED = 999
_C.MODE = "train"  # 运行模式选择 train/test/predict
_C.scriptPath = ""
"""
模型所需参数
"""
_C.MODEL = CN()
_C.MODEL.ARCH = "PSTA"  # 重识别模型选择
_C.MODEL.PRETRAIN_CHOICE = "imagenet"  # 使用基于imagenet训练出来的网络作为基础网络
_C.MODEL.NAME = "resnet50"
_C.MODEL.SEQ_LEN = 8
"""
DATALOADER参数
"""
_C.DATALOADER = CN()
_C.DATALOADER.NUM_WORKERS = 4
_C.DATALOADER.SAMPLER = 'softmax'
_C.DATALOADER.NUM_INSTANCE = 2
# ------------------------------------------------------------------------------------
# INPUT
# ------------------------------------------------------------------------------------
_C.INPUT = CN()
_C.INPUT.SIZE_TRAIN = [256, 128]
_C.INPUT.SIZE_TEST = [256, 128]
_C.INPUT.PROB = 0.5  # Random probability for image horizontal filp
_C.INPUT.RE_PROB = 0.5  # Random probability for random erasing
_C.INPUT.PIXEL_MEAN = [0.485, 0.456, 0.406]
_C.INPUT.PIXEL_STD = [0.229, 0.224, 0.225]
_C.INPUT.PADDING = 10  # Value of padding size
"""
主要用于设定远望定制的参数
"""
_C.params = CN()
"""
训练所需参数
"""
_C.params.CHECKPOINT_SAVE_PATH = ""  # 训练权重文件存储位置，每次训练前会自动清空该文件夹
_C.params.CHECKPOINT_SAVE_EPOCHS = 5  # 几个epoch保存一次预训练模型，将模型参数保存为epoch_num.pth
_C.params.TRAIN_IMAGES_PATH = ""
_C.params.TRAIN_FILE_PATH = ""
_C.params.BATCH_SIZE = 128
_C.params.WORKERS = 4
_C.params.START_EPOCH = 0  # 开始训练轮次
_C.params.VAL_EPOCHS = 5  # 模型验证频率
_C.params.VAL_START_EPOCHS = 50  # 开始验证的轮次
_C.params.MAX_EPOCHS = 50  # 最大训练轮次
_C.params.LOG_PATH = ""
_C.params.RESUME = False
_C.params.RESUME_WEIGHT_PATH = None
_C.params.ONNX_SAVE_PATH = ""
_C.params.TRAIN_SAMPLER = "Random_interval"  # 训练集采样策略 Random_interval / Begin_interval / dense
_C.params.TRIPLET_DISTANCE = "cosine"  # 三元组损失距离度量方法 cosine / euclidean

"""
训练优化器参数
"""
_C.SOLVER = CN()
_C.SOLVER.OPTIMIZER_NAME = "Adam"
# Base learning rate
_C.SOLVER.BASE_LR = 3e-4
# Factor of learning bias
_C.SOLVER.BIAS_LR_FACTOR = 2
_C.SOLVER.MOMENTUM = 0.9
_C.SOLVER.MARGIN = 0.3  # Margin of triplet loss
_C.SOLVER.CLUSTER_MARGIN = 0.3  # Margin of cluster

_C.SOLVER.CENTER_ON = 0
_C.SOLVER.CENTER_LR = 0.5
_C.SOLVER.CENTER_LOSS_WEIGHT = 0.0005

_C.SOLVER.RANGE_K = 2
_C.SOLVER.RANGE_MARGIN = 0.3
_C.SOLVER.RANGE_ALPHA = 0
_C.SOLVER.RANGE_BETA = 1
_C.SOLVER.RANGE_LOSS_WEIGHT = 1

# Setting of weight decay
_C.SOLVER.WEIGHT_DECAY = 0.0005
_C.SOLVER.WEIGHT_DECAY_BIAS = 0.

# decay rate of learning rate
_C.SOLVER.GAMMA = 0.1
# decay step of learning rate
_C.SOLVER.STEPS = (30, 55)

# warm up factor
_C.SOLVER.WARMUP_FACTOR = 1.0 / 3
# iterations of warm up
_C.SOLVER.WARMUP_ITERS = 500
# method of warm up , option : 'constant','linear'
_C.SOLVER.WARMUP_METHOD = "linear"

# epoch number of saving checkpoints
_C.SOLVER.CHECKPOINT_PERIOD = 50
# iteration of display training log
_C.SOLVER.LOG_PERIOD = 100
# epoch number of validation
_C.SOLVER.EVAL_PERIOD = 50

_C.SOLVER.FP_16 = True
_C.SOLVER.SEQS_PER_BATCH = 16
"""
测试所需参数
"""
_C.params.TEST_IMAGES_PATH = ""
_C.params.TEST_FILE_PATH = ""
_C.params.MODEL_PATH = ""
_C.params.RESULT_PATH = ""
_C.params.TEST_SAMPLER = "Begin_interval"  # 训练集采样策略 Random_interval / Begin_interval / dense
_C.params.TEST_MAX_SEQ_NUM = 200
_C.params.SEQS_PER_BATCH = 32
_C.params.TEST_DISTANCE = "cosine"  # 三元组损失距离度量方法 cosine / euclidean
"""
构建基础库/动态指纹库/重识别应用所需参数
"""
_C.params.BASELIB_FILE_PATH = ""
